import sys

from thrift.protocol import TCompactProtocol
from thrift.transport import TSocket, TTransport
from thrift_json import thrift2json

import config
from rpc import VulnQueryService
from rpc.ttypes import AsyncQueryResult, QueryHistoryResult
from tools.codec.factory import CodecFactory, Codec
from tools.file_helper import get_files, process_filepath


def parse_args():
    import argparse
    argument_parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description='The IoT firmwares vuln detection rpc-cli.',
        epilog="""
If you wanna submit a query request, you should do just like below:
    [E.g.]
        YOU WANT: scan all files under this dir.
            python3 cli.py [-H server-host -p server-port] -d my_firmware/
        YOU WANT: scan just some files.
            python3 cli.py [-H server-host -p server-port] -f my_firmware/file1.c my_firmware/file2.c
    And then you would get return immediately, this result just tells you the query key in this turn,
    the return just likes {"status_code": 0, "key": "cf109fbe-98c8-47cd-9ab2-7e0d55250847", "err_msg": ""}.
    
    ALL FIELDS IN RETURN:
        status_code:
            0: means OK
            -1: means some errors happened during the submission.
        key:
            an UUID string indicates an unique query key, used to query request result later.
        err_msg:
            an error message would be provided just when status_code=-1
If you wanna query the detection result to the request you submitted before, you should do just like below:
    [E.g.]
        python3 cli.py [-H server-host -p server-port] -q -k cf109fbe-98c8-47cd-9ab2-7e0d55250847
    ALL FIELDS IN RETURN:
        status_code:
            -1: means some errors happened during the vuln detection.
            0: means the query request you submitted before is now queuing.
            1: means the query request you submitted before is now running.
            2: means the query request you submitted before is done.
        progress:
            a float value(0~100), means the execution progress of vuln detection.
        vuln_results:
            a list, means detection result, would be supplied when status_code=2
        err_msg:
            an error massage would be provided just when status_code=-1
        queuing_cost:
            an int value, means total cost time in queuing.
        running_cost:
            an int value, means total cost time in running.
        """
    )
    argument_parser.add_argument('-f', '--input_files', dest='input_files', metavar='INPUT_FILE', nargs='*',
                                 help='Specifies which files will be used to extract.')
    argument_parser.add_argument('-d', '--input_dir', dest='input_dir', metavar='INPUT_DIR', nargs='?',
                                 help='Specifies which dir will be used to extract.')
    argument_parser.add_argument('-H', '--host', dest='host', default=config.RPC_SERVER_HOST, type=str,
                                 help='Specifies the server host ip.')
    argument_parser.add_argument('-p', '--port', dest='port', default=config.RPC_SERVER_PORT, type=int,
                                 help='Specifies the server port.')
    argument_parser.add_argument('-q', '--query', dest='is_query', action='store_true',
                                 help='Specifies if this action is a query.')
    argument_parser.add_argument('-k', '--key', dest='query_key', default='', type=str,
                                 help='Specifies query key.')
    argument_parser.add_argument('-o', '--out_type', dest='out_type', default=1, type=int, choices=(1, 2),
                                 help='Specifies output type, 1 means stdout and 2 means stderr.')

    def check(args):
        if not args.is_query:
            if not args.input_files and not args.input_dir:
                raise RuntimeError('Plz specify one of options -f or -d when not specifies -q.')
        else:
            if not args.query_key:
                raise RuntimeError('Plz specify a query key by -k when specifies -q.')
        return args

    return check(argument_parser.parse_args())


def handle(files: list):
    names = []
    datas = []
    for file in files:
        with open(file, 'rb') as f:
            data = codec.compress(f.read())
        names.append(process_filepath(file))
        datas.append(data)
    return names, datas


if __name__ == '__main__':
    args = parse_args()
    out = sys.stdout if args.out_type == 1 else sys.stderr

    socket = TSocket.TSocket(args.host, args.port)
    transport = TTransport.TFramedTransport(socket)
    protocol = TCompactProtocol.TCompactProtocol(transport)
    client = VulnQueryService.Client(protocol)

    codec: Codec = CodecFactory.new_gzip_codec()
    try:
        transport.open()

        if args.is_query:
            key = args.query_key
            ret: QueryHistoryResult = client.query_history(key)
        else:
            input_files = get_files(args.input_files, args.input_dir, config.VALID_SOURCE_SUFFICES)
            ret: AsyncQueryResult = client.async_query(*handle(input_files))
        print(thrift2json(ret), file=out)
    finally:
        transport.close()
